package main

import (
	"ai-gateway/common"
	"ai-gateway/config"
	_ "ai-gateway/config"
	"ai-gateway/db"
	"ai-gateway/proto"
	"context"
	"fmt"
	"net"

	"google.golang.org/grpc"
)

type TrafficServer struct {
	proto.UnimplementedTrafficServiceServer
}

// 检查业务系统的每日调用量是否超过配额（查quotas表和call_logs表）
func (s *TrafficServer) CheckQuota(ctx context.Context, req *proto.AIRequest) (*proto.QuotaResponse, error) {
	//查quotas表，获取该业务系统的每日配额
	var dailyQuota int
	err := db.DB.Get(&dailyQuota, "SELECT daily_quota FROM quotas WHERE business_system = ?", req.BusinessSystem)
	if err != nil {
		common.Logger.Error().Err(err).Msg("查询配额失败")
		return &proto.QuotaResponse{Allowed: false, Message: "配额配置错误"}, nil
	}
	//查call_logs表，统计今日已调用次数
	var used int
	err = db.DB.Get(&used, `
        SELECT COUNT(*) FROM call_logs 
        WHERE business_system = ? AND DATE(created_at) = CURDATE()
    `, req.BusinessSystem)
	if err != nil {
		common.Logger.Error().Err(err).Msg("统计已用次数失败")
		return &proto.QuotaResponse{Allowed: false, Message: "统计错误"}, nil
	}
	//判断是否超限
	allowed := used < dailyQuota
	msg := fmt.Sprintf("允许调用（已用 %d/%d）", used, dailyQuota)
	if !allowed {
		msg = fmt.Sprintf("今日配额已用尽（已用 %d/%d）", used, dailyQuota)
	}

	common.Logger.Info().
		Str("business_system", req.BusinessSystem).
		Bool("allowed", allowed).
		Msg("配额校验结果")

	return &proto.QuotaResponse{Allowed: allowed, Message: msg}, nil
}

func main() {
	common.InitLogger("traffic-service")
	if err := db.InitDB(); err != nil {
		common.Logger.Fatal().Err(err).Msg("数据库初始化失败")
	}

	lis, err := net.Listen("tcp", fmt.Sprintf("%s%d", ":", config.Config.TrafficServicePort))
	if err != nil {
		common.Logger.Fatal().Err(err).Msg("监听端口失败")
	}

	s := grpc.NewServer()
	proto.RegisterTrafficServiceServer(s, &TrafficServer{})
	common.Logger.Info().Msg("Traffic Service 启动: :8086")

	if err := s.Serve(lis); err != nil {
		common.Logger.Fatal().Err(err).Msg("服务启动失败")
	}
}
